import numpy as np
import time
def get_rotation(alpha, beta, gamma):
    R_alpha = np.array([
        [1, 0, 0],
        [0, np.cos(alpha), -np.sin(alpha)],
        [0, np.sin(alpha), np.cos(alpha)]
    ])

    R_beta = np.array([
        [np.cos(beta), 0, np.sin(beta)],
        [0, 1, 0],
        [-np.sin(beta), 0, np.cos(beta)]
    ])

    R_gamma = np.array([
        [np.cos(gamma), -np.sin(gamma), 0],
        [np.sin(gamma), np.cos(gamma), 0],
        [0, 0, 1]
    ])
    
    R = np.dot(R_gamma, np.dot(R_beta, R_alpha))
    return R

def get_measurement(t, orient, anchors, tags):
    r = []
    R = get_rotation(orient[0], orient[1], orient[2])
    for i in range(tags.shape[0]):
        for j in range(anchors.shape[0]):
            r_ij = anchors[j,:] - np.dot(R, tags[i,:].T).T - t
            r.append(np.linalg.norm(r_ij))
    return np.asarray(r)


def get_jacobian(anchors, tags, t, orient):
    anchors = anchors.reshape(-1,3)
    # sender.send(anchors)
    tags = tags.reshape(-1,3)
    t = t.reshape(-1,3)
    orient = orient.reshape(-1,3)
    alpha, beta, gamma = orient[0,0], orient[0,1], orient[0,2]
    
    # start_time = time.perf_counter()
    R_alpha = np.array([
        [1, 0, 0],
        [0, np.cos(alpha), -np.sin(alpha)],
        [0, np.sin(alpha), np.cos(alpha)]
    ])
    R_beta = np.array([
        [np.cos(beta), 0, np.sin(beta)],
        [0, 1, 0],
        [-np.sin(beta), 0, np.cos(beta)]
    ])
    R_gamma = np.array([
        [np.cos(gamma), -np.sin(gamma), 0],
        [np.sin(gamma), np.cos(gamma), 0],
        [0, 0, 1]
    ])
    # Define the time derivative of R_gamma
    dot_R_gamma = np.array([
        [-np.sin(gamma), -np.cos(gamma), 0],
        [np.cos(gamma), -np.sin(gamma), 0],
        [0, 0, 0]
    ])
    # Compute the time derivative of the rotation matrix R
    dot_R = np.dot(dot_R_gamma, np.dot(R_beta, R_alpha))
    R = np.dot(R_gamma, np.dot(R_beta, R_alpha))
    # dot_R = dot_R_gamma
    # R = np.eye(3)
    # print('rotation matrix time: ', time.perf_counter() - start_time)

    # start_time = time.perf_counter()
    # J1 = np.zeros((anchors.shape[0] * tags.shape[0], 4))
    # for i in range(tags.shape[0]):
    #     for j in range(anchors.shape[0]):
    #         r_ij = anchors[j,:] - np.dot(R, tags[i,:].T).T - t
    #         n_ij = r_ij / np.linalg.norm(r_ij)
    #         J1[i*anchors.shape[0]+j,0:3] = -n_ij
    #         J1[i*anchors.shape[0]+j,3:4] = -np.dot(n_ij, np.dot(dot_R, tags[i,:].T))
    # print('Jacobian 1 time: ', time.perf_counter() - start_time)

    
    num_anchors = anchors.shape[0]
    num_tags = tags.shape[0]
    J = np.zeros((num_anchors * num_tags, 4))
    R_tags =  np.dot(R, tags.T).T + t

    # start_time = time.perf_counter()
    for i in range(num_tags): 
        # 计算 r_ij 和 n_ij
        r_ij = anchors - R_tags[i]
        norms = np.linalg.norm(r_ij, axis=1)  # 计算每个 r_ij 的模长
        n_ij = r_ij / norms[:, None]  # 归一化 r_ij
        # 更新 J 矩阵的前 3 列
        J[i*num_anchors:(i+1)*num_anchors, 0:3] = -n_ij
        # 计算对偏移的导数
        dot_n_ij = -np.dot(n_ij, np.dot(dot_R, tags[i,:].T))
        # 更新 J 矩阵的最后一列
        J[i*num_anchors:(i+1)*num_anchors, 3:4] = dot_n_ij[:, None]
    # print('Jacobian 2 time: ', time.perf_counter() - start_time)

    # start_time = time.perf_counter()
    # r_ij = anchors[:, None, :] - R_tags  # [num_anchors, num_tags, 3]
    # norms = np.linalg.norm(r_ij, axis=2)  # 计算每个 r_ij 的模长 [num_anchors, num_tags]
    # n_ij = r_ij / norms[:, :, None]  # 归一化 r_ij [num_anchors, num_tags, 3]
    # # 更新 J 的前 3 列
    # J[:, 0:3] = -n_ij.reshape(-1, 3)  # [num_anchors * num_tags, 3]
    # # 计算 dot_n_ij 并更新 J 的第 4 列
    # dot_R_r_n = np.dot(dot_R, tags.T)
    # # dot_n_ij = np.einsum('ijk,kl->il', n_ij, np.dot(dot_R, tags.T))  # [num_anchors, num_tags, 3]
    # # J[:, 3:4] = dot_n_ij.reshape(-1, 3)[:, 0:1]  # [num_anchors * num_tags, 1]
    # J[0::2, 3:4] = np.dot(n_ij[:,0::2,:] ,dot_R_r_n[:,0])
    # J[1::2, 3:4] = np.dot(n_ij[:,1::2,:] ,dot_R_r_n[:,1])
    
    # print('Jacobian 3 time: ', time.perf_counter() - start_time)

    # with open(os.path.join(current_dir,"jacobians.txt"), "a") as f:
    #         f.write("J1:\n")
    #         np.savetxt(f, J1, fmt="%.6e", delimiter=",")
    #         f.write("\n")
    #         f.write("J:\n")
    #         np.savetxt(f, J, fmt="%.6e", delimiter=",")
    #         f.write("-------------------------------\n")

    return J
    

def get_crlb(anchors, tags, state):
    J = get_jacobian(anchors, tags, state[0:3], state[3:6])
    try:
        JTJ = np.matmul(J.transpose(), J)
        condition_number = np.linalg.cond(JTJ)  # 计算条件数
        # if condition_number > 1 / np.finfo(JTJ.dtype).eps:
        if condition_number > 1e7:
            raise np.linalg.LinAlgError("Matrix is near singular or poorly conditioned.")
        # 如果矩阵不是奇异的，则计算逆矩阵
        H = np.linalg.inv(JTJ) * 1e-2
    except np.linalg.LinAlgError as e:
        print(f"Error: {e}")
        # H = np.linalg.pinv(JTJ)  # 伪逆作为替代
        H = np.ones_like(JTJ) * 1e6
        print("Using large diagnal matrix instead.")
        # print(H)
    return H
